package com.atguigu.huffmancode;

import java.io.*;
import java.util.*;

public class HuffmanCodeDemo {
    public static void main(String[] args) {
//        String s="i like like like java do you like a java";
//        System.out.println("bytes:"+Arrays.toString(s.getBytes()));
//
//        HuffmanCode huffmanCode = HuffmanCode.getInstance(s);
//        // 压缩(编码)
//        byte[] zip = huffmanCode.zip(s);
//
//        System.out.println("zip:"+Arrays.toString(zip));
//
//
//        String unzip = huffmanCode.unzip2Str(zip);
//        System.out.println("unzip:"+unzip);

        //实测, 压缩后文件反而更大了. 推测只对文本有效

        FileZipUtil fileZipUtil = new FileZipUtil();

        String srcPath="D:\\src.png";
        String dstPath="D:\\huffman.zip";

        fileZipUtil.zipFile(srcPath,dstPath);
        System.out.println("压缩完成");

//        String srcPath="D:\\huffman.zip";
//        String dstPath="D:\\unzip.png";
//
//        fileZipUtil.unzipFile(srcPath,dstPath);
//        System.out.println("解压完成");

    }
}

class FileZipUtil{
    public void zipFile(String srcFile,String zipFile){
        // 创建文件的输入流, 文件输出流, 对象输出流
        InputStream is=null;
        OutputStream os=null;
        ObjectOutputStream oos=null;
        try {
            is=new FileInputStream(srcFile);
            byte[] bytes=new byte[is.available()];
            is.read(bytes);

            HuffmanCode huffmanCode = HuffmanCode.getInstance(bytes);
            byte[] zip = huffmanCode.zip(bytes);

            System.out.println("解压前:"+bytes.length+",解压后:"+zip.length+"+n*"+huffmanCode.getCodeMap().size());


            os=new FileOutputStream(zipFile);
            oos=new ObjectOutputStream(os);
            oos.writeObject(zip);
            oos.writeObject(huffmanCode.getCodeMap());

        }catch (IOException e){
            e.printStackTrace();
        }finally {
            try {
                oos.close();
                os.close();
                is.close();
            } catch (IOException e) {
                e.printStackTrace();
            }

        }
    }

    public void unzipFile(String zipFile,String unzipFile){
        InputStream is=null;
        ObjectInputStream ois=null;
        OutputStream os=null;

        try {
            is=new FileInputStream(zipFile);
            ois=new ObjectInputStream(is);
            byte[] bytes = (byte[]) ois.readObject();
            Map<Byte,String> codeMap=(Map<Byte, String>) ois.readObject();

            HuffmanCode huffmanCode = new HuffmanCode();
            huffmanCode.setCodeMap(codeMap);
            byte[] unzip = huffmanCode.unzip(bytes);

            os=new FileOutputStream(unzipFile);
            os.write(unzip);

        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } finally {
            try {
                os.close();
                ois.close();
                is.close();
            } catch (IOException e) {
                e.printStackTrace();
            }

        }
    }
}

class HuffmanCode {


    private Map<Byte,String> codeMap;

    public Map<Byte, String> getCodeMap() {
        return codeMap;
    }

    public void setCodeMap(Map<Byte, String> codeMap) {
        this.codeMap = codeMap;
    }

    private Node huffmanTree;

    public static HuffmanCode getInstance(String s){
        return getInstance(s.getBytes());
    }

    public static HuffmanCode getInstance(byte[] bytes){
        HuffmanCode huffmanCode = new HuffmanCode();
        huffmanCode.huffmanTree =huffmanCode.getHuffmanTree(bytes);
        huffmanCode.codeMap = huffmanCode.convertHuffmanTree2CodeMap(huffmanCode.huffmanTree);
        return huffmanCode;
    }

    public byte[] zip(String s){
        return zip(s.getBytes());
    }

    public byte[] zip(byte[] originBytes){
        //        System.out.println("code map:"+codeMap);
        String bitStr = this.convertStr2BitStrByHuffmanCodeMap(originBytes, codeMap);
//        System.out.println("bit str:"+bitStr);
        byte[] bytes = this.convertBitStr2Bytes(bitStr);
//        System.out.println("bytes:"+Arrays.toString(bytes));
        return bytes;
    }

    public String unzip2Str(String s){
        return unzip2Str(s.getBytes());
    }

    public String unzip2Str(byte[] originBytes){
        byte[] resultByteList=unzip(originBytes);

        // 3. 将byte集合转化为一个字符串
        StringBuilder resultSb=new StringBuilder();
        for (byte b : resultByteList) {
            resultSb.append((char)b);
        }
        return resultSb.toString();
    }

    public byte[] unzip(byte[] bytes){
        // 1. 将byte数组转化为bit字符串
        StringBuilder bitSb=new StringBuilder();
        for(int i=0;i<bytes.length;i++){
            String bitStr = convertByte2BitStr(bytes[i], i == bytes.length - 1);
            bitSb.append(bitStr);
        }

//        System.out.println("bit str:"+bitSb.toString());

        // 获取bit字符串转byte的map
        Map<String,Byte> bitStr2ByteMap=new HashMap<>();
        for (Map.Entry<Byte, String> entry : codeMap.entrySet()) {
            bitStr2ByteMap.put(entry.getValue(),entry.getKey());
        }
        System.out.println("查找速度:"+bitStr2ByteMap.size());

        // 2. 使用快慢指针, 将bitStr逐次转化为byte集合
        int slowIndex=0;
        List<Byte> resultByteList=new ArrayList<>();
        for(int fastIndex=1;fastIndex<bitSb.toString().length();fastIndex++){
            if(bitStr2ByteMap.containsKey(bitSb.substring(slowIndex,fastIndex))){
                resultByteList.add(bitStr2ByteMap.get(bitSb.substring(slowIndex,fastIndex)));
                slowIndex=fastIndex;
            }
            if(fastIndex%0x1000==0){
                System.out.println("解压进度:"+resultByteList.size());
            }
        }
//        System.out.println("result byte:"+resultByteList);

        byte[] resultBytes=new byte[resultByteList.size()];
        for (int i = 0; i < resultByteList.size(); i++) {
            resultBytes[i]=resultByteList.get(i);
        }
        return resultBytes;
    }



    /**
     * 获取huffman树
     * @return
     */
    public Node getHuffmanTree(byte[] bytes){
        if (bytes.length==0) {
            return null;
        }
        // 字符串转成byte数组

        // 构建字符权重节点map
        Map<Byte,Integer> byte2WeightMap=new HashMap<>();
        for (byte aByte : bytes) {
            byte2WeightMap.merge(aByte, 1, Integer::sum);
        }

        // 根据字符权重map构建节点集合
        List<Node> nodeList=new ArrayList<>(byte2WeightMap.size());
        for (Map.Entry<Byte,Integer> entry:byte2WeightMap.entrySet()) {
            nodeList.add(new Node(entry.getKey(), entry.getValue()));
        }

        // 根据字符权重集合进行排序, 逐次构建huffman树
        while(nodeList.size()>1){
            Collections.sort(nodeList);

            Node node1 = nodeList.get(0);
            Node node2 = nodeList.get(1);
            Node merge = Node.merge(node1, node2);

            nodeList.remove(node1);
            nodeList.remove(node2);
            nodeList.add(merge);
        }

        return nodeList.get(0);
    }

    /**
     * 获取huffman编码表
     * @param huffmanNode
     * @return
     */
    public Map<Byte,String> convertHuffmanTree2CodeMap(Node huffmanNode){
        Map<Byte,String> byte2HuffmanCodeMap=new HashMap<>();
        if (huffmanNode==null) {
            return byte2HuffmanCodeMap;
        }

        // 分别遍历至子节点(考虑使用递归实现)
        goNext(huffmanNode,byte2HuffmanCodeMap,"");

        return byte2HuffmanCodeMap;
    }

    /**
     * huffman编码表内部递归方法
     * @param node
     * @param byte2HuffmanCodeMap
     * @param parentPath
     */
    private void goNext(Node node,Map<Byte,String> byte2HuffmanCodeMap,String parentPath){
        assert node!=null;
        if(node.b!=null){
            byte2HuffmanCodeMap.put(node.b,parentPath);
            return;
        }
        if(node.left!=null){
            goNext(node.left,byte2HuffmanCodeMap,parentPath+"0");
        }
        if(node.right!=null){
            goNext(node.right,byte2HuffmanCodeMap,parentPath+"1");
        }

    }


    public String convertStr2BitStrByHuffmanCodeMap(byte[] bytes,Map<Byte,String> huffmanCodeMap){
        StringBuilder stringBuilder=new StringBuilder();
        for (byte b : bytes) {
            stringBuilder.append(huffmanCodeMap.get(b));
        }
        return stringBuilder.toString();
    }

    public byte[] convertBitStr2Bytes(String bitStr){
        int byteLen=(bitStr.length()-1)/8+1;
        byte[] bytes=new byte[byteLen];

        for(int i=0;i<bitStr.length();i+=8){
            String eachByteStr = (i+8>bitStr.length())?bitStr.substring(i):bitStr.substring(i, i + 8);

            bytes[i/8]=(byte)Integer.parseInt(eachByteStr,2);
        }

        return bytes;
    }

    public String convertByte2BitStr(byte b,boolean isEnd){
        int i = b;

        if(isEnd){
            return Integer.toBinaryString(i);
        }else{
            i|=0x100;
            String binaryString = Integer.toBinaryString(i);

            return binaryString.substring(binaryString.length() - 8);
        }
    }

}

class Node implements Comparable<Node>{
    public Byte b;
    public int weight;
    public Node left;
    public Node right;

    public Node(Byte b, int weight) {
        this.b = b;
        this.weight = weight;
    }


    public Node(Byte b, int weight, Node left, Node right) {
        this.b = b;
        this.weight = weight;
        this.left = left;
        this.right = right;
    }

    @Override
    public int compareTo(Node o) {
        return this.weight-o.weight;
    }

    public static Node merge(Node node1,Node node2){
        return new Node(null,node1.weight+node2.weight,node1,node2);
    }

    public void preOrder(){
        System.out.println(this);

        if(this.left!=null){
            this.left.preOrder();
        }

        if(this.right!=null){
            this.right.preOrder();
        }
    }

    @Override
    public String toString() {
        return "Node{" +
                "b=" + b +
                ", weight=" + weight +
                '}';
    }
}
